# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: reporter\reporter.py
import logging, json, traceback, os, os.path, datetime, re
from subprocess import check_call
from db import Connection
from sqlalchemy import select as sql_select
from sqlalchemy import distinct as sql_distinct
from db.tables.reports import ReportRow, ReportsTable
from db.tables.scans import ScanRow, ScansTable
from db.tables.scan_sessions import ScanSessionRow, ScanSessionsTable
from db.tables.report_templates import (
    ReportTemplateRow,
    ReportTemplatesTable,
    ReportType,
)
from db.tables.target2group import Target2GroupRow
from reporter.generators.ds_scan import DSScanSessionResult
from reporter.generators.ds_target import DSTargetResult
from db.tables.target_vulns import TargetVulnRow
from db.tables.scan_session_vulns import ScanSessionVulnRow
from api.classes.frontend.auto import *
from reporter.reports.quick import QuickReport
from reporter.reports.developer import DeveloperReport
from reporter.reports.compliance import ComplianceReport
from reporter.reports.affected_items import AffectedItemsReport
from reporter.reports.executive_summary import ExecutiveSummaryReport
from reporter.reports.exporter import Exporter
from reporter.reports.exporter_imperva import ExporterImperva
from reporter.reports.exporter_f5 import ExporterF5
from reporter.reports.scan_comparison import ScanComparison
from reporter.reports.exporter_fortinet_vulnerabilities import (
    ExporterFortinetVulnerabilities,
)
from db.tables.events import create_event, EventResources
from db.tables.users import UsersTable, UserRow

logger = logging.getLogger("reporter")
re_source_non_alphanumeric = re.compile("[^\\w]+")


class Reporter:
    def __init__(
        self, shard_id, reports_local_storage, pdf_converter=None, s3_storage=None
    ):
        self.shard_id = shard_id
        self.reports_storage = reports_local_storage
        self.pdf_converter = pdf_converter
        self.s3_storage = s3_storage

    def get_next_report(self):
        with Connection(self.shard_id) as (db_connection):
            select_from = ReportsTable.join(
                ReportTemplatesTable,
                ReportRow.report_template_id == ReportTemplateRow.report_template_id,
            )
            select_from = select_from.join(
                UsersTable, UserRow.user_id == ReportRow.creator_id
            )
            query = (
                sql_select(
                    (
                        ReportRow,
                        ReportTemplateRow.data,
                        ReportTemplateRow.type,
                        ReportTemplateRow.name,
                        UserRow.email.label("creator_email"),
                    )
                )
                .select_from(select_from)
                .where(ReportRow.status == "queued")
                .where(ReportRow.deleted_at.is_(None))
                .order_by(ReportRow.created)
                .limit(1)
            )
            report = db_connection.execute(query).fetchone()
            if not report:
                return
            return dict(report)

    def get_scan_data_source(self, scan_id, source_type=None):
        with Connection(self.shard_id) as (db_connection):
            select_from = ScansTable.join(
                ScanSessionsTable,
                ScanRow.current_session_id == ScanSessionRow.scan_session_id,
            )
            query = (
                sql_select((ScanRow, ScanSessionRow.allowed_targets))
                .select_from(select_from)
                .where(ScanRow.scan_id == scan_id)
            )
            scan = db_connection.execute(query).fetchone()
            if not scan:
                return
            current_session_id = scan.current_session_id
            allowed_targets = scan.allowed_targets
            scan_target_id = scan.target_id
        if not allowed_targets:
            yield DSScanSessionResult(
                self.shard_id, current_session_id, source_type=source_type
            )
        else:
            yield DSScanSessionResult(
                self.shard_id,
                current_session_id,
                filter_target_id=scan_target_id,
                source_type=source_type,
            )
            for host, tid in allowed_targets.items():
                yield DSScanSessionResult(
                    self.shard_id,
                    current_session_id,
                    filter_target_id=tid,
                    source_type=source_type,
                )

    def get_scan_session_data_source(
        self, scan_session_id, vuln_id_list=None, source_type=None
    ):
        query = sql_select(
            (ScanSessionRow.allowed_targets, ScanSessionRow.target_id)
        ).where(ScanSessionRow.scan_session_id == scan_session_id)
        with Connection(self.shard_id) as (db_connection):
            scan_session = db_connection.execute(query).fetchone()
            if not scan_session:
                return
            allowed_targets = scan_session.allowed_targets
            scan_session_target_id = scan_session.target_id
        if not allowed_targets:
            yield DSScanSessionResult(
                self.shard_id,
                scan_session_id,
                vuln_id_list=vuln_id_list,
                source_type=source_type,
            )
        else:
            yield DSScanSessionResult(
                self.shard_id,
                scan_session_id,
                filter_target_id=scan_session_target_id,
                vuln_id_list=vuln_id_list,
                source_type=source_type,
            )
            for host, target_id in allowed_targets.items():
                yield DSScanSessionResult(
                    self.shard_id,
                    scan_session_id,
                    filter_target_id=target_id,
                    vuln_id_list=vuln_id_list,
                    source_type=source_type,
                )

    def generate_data_sources(self, report):
        """
        Report sources:
            all_vulnerabilities,
            targets
            groups
            scans
            scan_result
            vulnerabilities
        :param report:
        :return:
        """
        source = ClsReportSource(**report["source"])
        if source.list_type == "all_vulnerabilities":
            ds_list = []
            for target_id in source.id_list:
                ds = DSTargetResult(
                    self.shard_id, target_id, source_type=source.list_type
                )
                ds_list.append(ds)

            return ds_list
        if source.list_type == "targets":
            ds_list = []
            for target_id in source.id_list:
                ds = DSTargetResult(
                    self.shard_id, target_id, source_type=source.list_type
                )
                ds_list.append(ds)

            return ds_list
        if source.list_type == "groups":
            ds_list = []
            query = sql_select((sql_distinct(Target2GroupRow.target_id),)).where(
                Target2GroupRow.group_id.in_(source.id_list)
            )
            target_id_list = set()
            with Connection(self.shard_id) as (db_connection):
                for target in db_connection.execute(query).fetchall():
                    target_id_list.add(target.target_id)

            for target_id in target_id_list:
                ds = DSTargetResult(
                    self.shard_id, target_id, source_type=source.list_type
                )
                ds_list.append(ds)

            return ds_list
        if source.list_type in ("scans", "scan_pair"):
            ds_list = []
            for scan_id in source.id_list:
                for ds in self.get_scan_data_source(
                    scan_id, source_type=source.list_type
                ):
                    ds_list.append(ds)

            return ds_list
        if source.list_type in ("scan_result", "scan_result_pair"):
            ds_list = []
            for scan_session_id in source.id_list:
                for ds in self.get_scan_session_data_source(
                    scan_session_id, source_type=source.list_type
                ):
                    ds_list.append(ds)

            return ds_list
        if source.list_type == "vulnerabilities":
            query = sql_select((sql_distinct(TargetVulnRow.target_id),)).where(
                TargetVulnRow.vuln_id.in_(source.id_list)
            )
            with Connection(self.shard_id) as (db_connection):
                target_id_list = set()
                for target in db_connection.execute(query).fetchall():
                    target_id_list.add(target.target_id)

            ds_list = []
            for target_id in target_id_list:
                ds = DSTargetResult(
                    self.shard_id,
                    target_id,
                    source.id_list,
                    source_type=source.list_type,
                )
                ds_list.append(ds)

            return ds_list
        if source.list_type == "scan_vulnerabilities":
            query = sql_select(
                (sql_distinct(ScanSessionVulnRow.scan_session_id),)
            ).where(ScanSessionVulnRow.vuln_id.in_(source.id_list))
            with Connection(self.shard_id) as (db_connection):
                scan_session_id_list = set()
                for scan_session in db_connection.execute(query).fetchall():
                    scan_session_id_list.add(scan_session.scan_session_id)

            ds_list = []
            for scan_session_id in scan_session_id_list:
                for ds in self.get_scan_session_data_source(
                    scan_session_id, source.id_list, source_type=source.list_type
                ):
                    ds_list.append(ds)

            return ds_list

    @staticmethod
    def get_reporter_params(report_data):
        reporter_init_params = dict()
        reporter_class_name = report_data["class"]
        if reporter_class_name == "DeveloperReport":
            reporter_class = DeveloperReport
        else:
            if reporter_class_name == "QuickReport":
                reporter_class = QuickReport
            else:
                if reporter_class_name == "ComplianceReport":
                    reporter_init_params["xml_name"] = report_data["xml"]
                    reporter_class = ComplianceReport
                else:
                    if reporter_class_name == "AffectedItemsReport":
                        reporter_class = AffectedItemsReport
                    else:
                        if reporter_class_name == "ExecutiveSummaryReport":
                            reporter_class = ExecutiveSummaryReport
                        else:
                            if reporter_class_name == "Exporter":
                                reporter_class = Exporter
                            else:
                                if reporter_class_name == "ExporterImperva":
                                    reporter_class = ExporterImperva
                                else:
                                    if reporter_class_name == "ExporterF5":
                                        reporter_class = ExporterF5
                                    else:
                                        if (
                                            reporter_class_name
                                            == "ExporterFortinetVulnerabilities"
                                        ):
                                            reporter_class = (
                                                ExporterFortinetVulnerabilities
                                            )
                                        else:
                                            if reporter_class_name == "ScanComparison":
                                                reporter_class = ScanComparison
                                            else:
                                                raise Exception(
                                                    "template not known %s",
                                                    reporter_class_name,
                                                )
                                            return (
                                                reporter_class,
                                                reporter_init_params,
                                            )

    def generate_report(self, report):
        with Connection(self.shard_id) as (db_connection):
            query = (
                (ReportsTable.update(values=dict(status="processing")))
                .where(ReportRow.report_id == report["report_id"])
                .where(ReportRow.status == "queued")
            )
            result = db_connection.execute(query)
            if result.rowcount != 1:
                logger.error("report not found aborting %s", result)
                return
        today = datetime.date.today()
        source_descriptor = report["source"].get("description", "")
        if ";" in report["source"].get("description"):
            descriptors = []
            for source in source_descriptor.split(","):
                descriptors.append(
                    re_source_non_alphanumeric.sub("_", source.split(";")[0])
                )

            source_descriptor = ("_").join(descriptors)
        report_file_name_nice_wo_extension = ("_").join(
            (
                today.strftime("%Y%m%d"),
                re_source_non_alphanumeric.sub("_", report["name"]).replace(" ", "_"),
                source_descriptor.replace(" ", "_"),
            )
        )
        event_data = dict(
            name=None,
            owner_id=report["owner_id"],
            user_id=report["creator_id"],
            data=dict(
                creator=report["creator_email"],
                name=report_file_name_nice_wo_extension,
                template_name=report["name"],
                source=report["source"].get("description"),
                source_type=report["source"].get("list_type"),
            ),
            resource_type=EventResources.export,
            resource_id=report["report_id"],
            shard=self.shard_id,
        )
        try:
            reporter_class, reporter_init_params = self.get_reporter_params(
                report["data"]
            )
            ds = self.generate_data_sources(report)
            if not ds:
                raise Exception("no data source found")
            reporter_init_params["sources"] = ds
            reporter = reporter_class(**reporter_init_params)
            generated_files = []
            if report["type"] == ReportType.report:
                html_report_file_mame = report["report_id"] + ".html"
                html_report_file_path = os.path.join(
                    self.reports_storage, html_report_file_mame
                )
                reporter.generate_report(html_report_file_path)
                if not self.s3_storage:
                    generated_files.append(html_report_file_mame)
                    open(
                        os.path.join(
                            self.reports_storage, html_report_file_mame + ".info"
                        ),
                        "wt",
                    ).write(
                        json.dumps(
                            {"name": report_file_name_nice_wo_extension + ".html"}
                        )
                    )
                else:
                    s3_file_name = "%s/%s/%s.html" % (
                        report["owner_id"],
                        report["report_id"],
                        report_file_name_nice_wo_extension,
                    )
                    generated_files.append(s3_file_name)
                    from helpers.aws.s3 import upload_file

                    upload_file(
                        self.s3_storage["bucket"],
                        self.s3_storage["dir"] + "/" + s3_file_name,
                        html_report_file_path,
                        content_disposition='attachment; filename="%s%s"'
                        % (report_file_name_nice_wo_extension, ".html"),
                    )
                if self.pdf_converter:
                    pdf_file_path = os.path.join(
                        self.reports_storage, report["report_id"] + ".pdf"
                    )
                    check_call(
                        [self.pdf_converter, html_report_file_path, pdf_file_path]
                    )
                    pdf_report_file_mame = report["report_id"] + ".pdf"
                    if not self.s3_storage:
                        generated_files.append(pdf_report_file_mame)
                        open(
                            os.path.join(
                                self.reports_storage, pdf_report_file_mame + ".info"
                            ),
                            "wt",
                        ).write(
                            json.dumps(
                                {"name": report_file_name_nice_wo_extension + ".pdf"}
                            )
                        )
                    else:
                        from helpers.aws.s3 import upload_file

                        s3_file_name = "%s/%s/%s.pdf" % (
                            report["owner_id"],
                            report["report_id"],
                            report_file_name_nice_wo_extension,
                        )
                        generated_files.append(s3_file_name)
                        upload_file(
                            self.s3_storage["bucket"],
                            self.s3_storage["dir"] + "/" + s3_file_name,
                            pdf_file_path,
                            content_disposition='attachment; filename="%s%s"'
                            % (report_file_name_nice_wo_extension, ".pdf"),
                        )
                        os.remove(html_report_file_path)
                        os.remove(pdf_file_path)
                else:
                    if report["type"] == ReportType.export:
                        ext = report["data"].get("ext", ".xml")
                        report_file_mame = report["report_id"] + ext
                        html_report_file_path = os.path.join(
                            self.reports_storage, report_file_mame
                        )
                        reporter.generate_report(html_report_file_path)
                        if not self.s3_storage:
                            generated_files = [report_file_mame]
                            open(
                                os.path.join(
                                    self.reports_storage, report_file_mame + ".info"
                                ),
                                "wt",
                            ).write(
                                json.dumps(
                                    {"name": report_file_name_nice_wo_extension + ext}
                                )
                            )
                        else:
                            from helpers.aws.s3 import upload_file

                            s3_file_name = "%s/%s/%s%s" % (
                                report["owner_id"],
                                report["report_id"],
                                report_file_name_nice_wo_extension,
                                ext,
                            )
                            generated_files.append(s3_file_name)
                            upload_file(
                                self.s3_storage["bucket"],
                                self.s3_storage["dir"] + "/" + s3_file_name,
                                html_report_file_path,
                                content_type="application/octet-stream",
                                content_disposition='attachment; filename="%s%s"'
                                % (report_file_name_nice_wo_extension, ext),
                            )
                            os.remove(html_report_file_path)
                        with Connection(self.shard_id) as (db_connection):
                            query = (
                                ReportsTable.update(
                                    values=dict(
                                        status="completed",
                                        result=dict(files=generated_files),
                                    )
                                )
                            ).where(ReportRow.report_id == report["report_id"])
                            db_connection.execute(query)
                            if report["type"] == ReportType.export:
                                event_data["name"] = "export_created"
                                event_data["resource_type"] = EventResources.export
                            else:
                                if report["type"] == ReportType.report:
                                    event_data["name"] = "report_created"
                                    event_data["resource_type"] = EventResources.report
                            query = create_event(**event_data)
                            db_connection.execute(query)
                        logger.debug("report %s done", report)
        except Exception as e:
            logger.exception("report %s failed %s", report, e)
            with Connection(self.shard_id) as (db_connection):
                query = (
                    ReportsTable.update(
                        values=dict(
                            status="failed",
                            result=dict(
                                traceback=traceback.format_exc(), exception=str(e)
                            ),
                        )
                    )
                ).where(ReportRow.report_id == report["report_id"])
                db_connection.execute(query)
                if report["type"] == ReportType.export:
                    event_data["name"] = "export_failed"
                    event_data["resource_type"] = EventResources.export
                else:
                    if report["type"] == ReportType.report:
                        event_data["name"] = "report_failed"
                        event_data["resource_type"] = EventResources.report
                query = create_event(**event_data)
                db_connection.execute(query)
